import json
import os

from PIL import Image
import cv2


def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def get_data_list(labels_file, cls_name):
    with open(labels_file, "r") as fp:
        data_json = json.load(fp)

    categories = data_json["categories"]
    annotations = data_json["annotations"]
    images = data_json["images"]

    categories_dic = {}
    images_dic = {}

    categorie_id = 0
    for categorie in categories:
        categories_dic[categorie["id"]] = categorie_id
        categorie_id += 1

    image_id = 0
    for image in images:
        images_dic[image["id"]] = image_id
        image_id += 1

    data_list = []

    for annotation in annotations:

        data_dic = {"cls_name": categories[categories_dic[annotation["category_id"]]]["name"],
                    "category_id": categories[categories_dic[annotation["category_id"]]]["id"],
                    "cls_serial_number": categories_dic[annotation["category_id"]],
                    "supercategory": categories[categories_dic[annotation["category_id"]]]["supercategory"],
                    "segmentation": annotation["segmentation"], "bbox": annotation["bbox"],
                    "image_id": annotation["image_id"],
                    "image_name": images[images_dic[annotation["image_id"]]]["file_name"],
                    "width": images[images_dic[annotation["image_id"]]]["width"],
                    "height": images[images_dic[annotation["image_id"]]]["height"]}

        if data_dic["cls_name"] == cls_name or cls_name is None:
            data_list.append(data_dic)

    return data_list


class cocoAPI:
    def __init__(self, labels_file, cls_name=None):
        self.data_list = get_data_list(labels_file, cls_name)

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, item):
        img_name = self.data_list[item][0]
        img = Image.open(img_name, "r")
        if len(img.split()) != 3:
            img = img.convert("RGB")
        targets = self.data_list[item][1]
        return img, targets

    def get_list(self):
        return self.data_list

    def to_VOC(self, save_path, images_path):
        if not os.path.exists(save_path + "/labels"):
            os.makedirs(save_path + "/labels")
        if not os.path.exists(save_path + "/JPEGImages"):
            os.makedirs(save_path + "/JPEGImages")

        for index in self.data_list:
            cls_id = index["cls_serial_number"]
            x_min, y_min, w, h = index["bbox"]
            image_name = index["image_name"]
            width = index["width"]
            height = index["height"]

            image = cv2.imread(images_path + "/" + image_name)
            cv2.imwrite(save_path + "/JPEGImages/" + image_name, image)

            bbox = convert((width, height), (x_min, x_min + w, y_min, y_min + h))

            out_file = open(save_path + "/labels/%s.txt" % image_name.split(".")[0], 'a')
            out_file.write(str(cls_id) + " " + " ".join([str(b) for b in bbox]) + '\n')
            out_file.close()


if __name__ == "__main__":
    data = cocoAPI("/home/anmc/data/annotations/instances_train2014.json")
    data.to_VOC("VOC", "/home/anmc/data/train2014/")
    print(len(data))
